# -*- coding: utf-8 -*-
# @Time : 2023/5/10 16:20
# @Author : 陈鹏飞
# @Email ： 2578925789@qq.com
# @File : analyseDataset
# @Description : 分析COCO数据集中各类别标签的分布情况，包括类别分布、尺寸分布等

import os
import json
import numpy as np
from tqdm import tqdm

if __name__ == "__main__":
    dataset_json = "./save/val.json"
    # dataset_json = "C:/Users/Administrator/Desktop/1/train_v1_resample3.json"
    model_h = 192
    model_w = 320
    print("******model_w:{} model_h:{}******".format(int(model_w), int(model_h)))
    class_item_dict = {}
    with open(dataset_json, "r", encoding="utf-8") as f:
        content = json.load(f)
        images_dict = {}
        for i, imginfo in enumerate(content["images"]):
            images_dict[imginfo["id"]] = imginfo
        categories_dict = {}
        for i, cateinfo in enumerate(content["categories"]):
            categories_dict[cateinfo["id"]] = cateinfo["name"]
            class_item_dict[cateinfo["name"]] = []
        for i, ann in enumerate(tqdm(content["annotations"])):
            imginfo = list(images_dict.values())[list(images_dict.keys()).index(ann["image_id"])]
            imgH = imginfo["height"]
            imgW = imginfo["width"]
            ratio = (float)(imgW) / model_w
            item_resize_w = (float)(ann["bbox"][2]) / ratio
            item_resize_h = (float)(ann["bbox"][3]) / ratio
            class_name = categories_dict[ann["category_id"]]
            class_item_dict[class_name].append([item_resize_w, item_resize_h])

        for k,v in class_item_dict.items():
            print("--------------{}--------------".format(k))
            print("nums:{} prop:{:.1%}".format(len(v), len(v) / len(content["annotations"])))
            item_singleC = np.array(v)
            ave_w, ave_h = np.average(item_singleC, axis=0)
            print("ave_w:{} ave_h:{}".format(int(ave_w), int(ave_h)))
            area = item_singleC[:, 0] * item_singleC[:, 1]
            max_index = np.argmax(area)
            max_w = item_singleC[max_index][0]
            max_h = item_singleC[max_index][1]
            print("max_w:{} max_h:{}".format(int(max_w), int(max_h)))
            min_index = np.argmin(area)
            min_w = item_singleC[min_index][0]
            min_h = item_singleC[min_index][1]
            print("min_w:{} min_h:{}".format(int(min_w), int(min_h)))






